#!/usr/bin/env python
# -*- coding:utf-8 _*-
# LangChain:万能的非结构化文档载入详解：
# https://zhuanlan.zhihu.com/p/624812261
# https://integrations.langchain.com/

# sudo apt-get update
# sudo apt-get install tesseract-ocr
# sudo apt-get install tesseract-ocr-chi-sim  # 安装简体中文语言包
# sudo apt-get install tesseract-ocr-chi-tra  # 安装繁体中文语言包
# pip install pytesseract

import os, shutil, platform, multiprocessing, time, uuid, nltk, uvicorn, json, math
# nltk.download('punkt')
# exit()

from tqdm import tqdm
from datetime import datetime
from typing import  List
from fastapi import FastAPI, Request, UploadFile, File
from options import config
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.document_loaders import (PyPDFLoader, BSHTMLLoader, TextLoader, CSVLoader, PythonLoader, JSONLoader,
                                        UnstructuredMarkdownLoader, UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader, UnstructuredExcelLoader, UnstructuredImageLoader, UnstructuredFileLoader,
                                        PDFMinerLoader, Docx2txtLoader)
# from langchain.document_loaders.image import UnstructuredImageLoader
from langchain.text_splitter import CharacterTextSplitter

embeddings = None

filetype_load_map = {
    'txt': TextLoader,
    'md': TextLoader,
    'py': PythonLoader,
    'json': JSONLoader,
    'pdf': PyPDFLoader, #PDFMinerLoader,
    'csv': CSVLoader,
    'html': BSHTMLLoader,
    'md': UnstructuredMarkdownLoader,
    'doc': UnstructuredFileLoader,
    'docx': Docx2txtLoader,
    'ppt': UnstructuredPowerPointLoader,
    'pptx': UnstructuredPowerPointLoader,
    'xls': UnstructuredExcelLoader,
    'xlsx': UnstructuredExcelLoader,
    'jpg': UnstructuredImageLoader,
    'jpeg': UnstructuredImageLoader,
    'png': UnstructuredImageLoader
}

def init():
    if os.path.exists(config.cache_path) == False:
        os.mkdir(config.cache_path)

app = FastAPI()
def loadEmbeddings():
    global embeddings
    if embeddings == None:
        print('载入Embeddings模型...')
        embeddings = HuggingFaceEmbeddings(model_name=config.embedding_model)

@app.post("/remove")
async def chatForKnowledge(request: Request):
    global embeddings
    now = datetime.now()
    Time = now.strftime("%Y-%m-%d %H:%M:%S")
    args = await request.form()
    if args['name'] == None or args['name'] == '':
        return {
            "response": "projectName is empty",
            "status": 500,
            "time": Time
        }
    projectName = ''
    try:
        projectName = str(args['name'] or '').strip()
    except:
        ...
    size = ''
    try:
        size = str(args['size'] or 2000).strip()
    except:
        ...
    path = os.path.join(config.wiki_path, projectName)
    if size != '':
        path = os.path.join(path, size)

    if os.path.exists(path):
        shutil.rmtree(path)
    return {
        "response": 'ok',
        "status": 200,
        "time": Time
    }

@app.post("/read")
async def chatForKnowledge(request: Request):
    global embeddings
    now = datetime.now()
    Time = now.strftime("%Y-%m-%d %H:%M:%S")
    args = await request.form()
    if args['name'] == None or args['name'] == '':
        return {
            "response": "projectName is empty",
            "status": 500,
            "time": Time
        }
    if args['query'] == None or args['query'] == '':
        return {
            "response": "query is empty",
            "status": 500,
            "time": Time
        }
    projectName = ''
    try:
        projectName = str(args['name'] or '').strip()
    except:
        ...
    query = ''
    try:
        query = str(args['query'] or '').strip()
    except:
        ...
    size = ''
    try:
        size = str(args['size'] or 2000).strip()
    except:
        ...
    k = args['k']
    if k == None: k = 2
    k = int(k)
    if math.isnan(k) == True: k = 2
    path = os.path.join(config.wiki_path, projectName)
    if os.path.exists(path) == False:
        return {
            "response": "The selected knowledge does not exist",
            "status": 500,
            "time": Time
        }
    
    if size == '':
        file_list = os.listdir(path)
        dir_list = [item for item in file_list if os.path.isdir(os.path.join(path, item))]
        if len(dir_list) > 0:
            size = dir_list[0].strip()

    if size != '':
        path = os.path.join(path, size)
        if os.path.exists(path) == False:
            return {
                "response": "The selected knowledge does not exist",
                "status": 500,
                "time": Time
            }
        
    loadEmbeddings()
    db = FAISS.load_local(path, embeddings)
    docs = db.similarity_search(query, k=2)
    return {
        "response": docs,
        "status": 200,
        "time": Time
    }

@app.post("/importByDir")
async def importByDir(request: Request):
    global embeddings
    now = datetime.now()
    Time = now.strftime("%Y-%m-%d %H:%M:%S")
    args = await request.form()
    params = []
    for item in request.query_params:
        params.append({
            item: str(request.query_params[item]).strip()
        })
    for item in args:
        params.append({
            item: str(args[item]).strip()
        })
    size = 2000
    try:
        size = int(args['size'])
    except:
        ...

    if args['path'] == None or args['path'] == '':
        return {
            "response": "path is empty",
            "status": 500,
            "time": Time
        }
    else:
        raw_texts = []
        metadatas = []
        path = str(args['path']).strip()
        root = os.listdir(path)
        dirs = [item for item in root if os.path.isdir(os.path.join(path, item))]
        for dir in dirs:  
            files = os.listdir(os.path.join(path, dir))
            for file in files:
                file = os.path.join(path, dir, file)
                if os.path.isdir(file) == False:
                    file_type = file.split('.')[-1].lower()
                    if file_type not in filetype_load_map:
                        ...
                    else:
                        projectName = str(dir).strip()
                        loadDocsFiles(file, size, projectName, metadatas, raw_texts, file_type)
                    ...
        return {
            "params": params,
            "response": dirs,
            "status": 200,
            "time": Time
        }
    
@app.post("/import")
async def importFile(request: Request, files: List[UploadFile] = File(...)):
    global embeddings
    args = await request.form()
    now = datetime.now()
    Time = now.strftime("%Y-%m-%d %H:%M:%S")
    params = []
    for item in request.query_params:
        params.append({
            item: str(request.query_params[item]).strip()
        })
    for item in args:
        params.append({
            item: str(args[item]).strip()
        })
    size = 2000
    try:
        size = int(args['size'])
    except:
        ...
    if not files or len(files) == 0:
        return {
            "response": "please upload your knowledge file",
            "status": 500,
            "time": Time
        }
    else:
        raw_texts = []
        metadatas = []
        for file in files:
            file_type = file.filename.split('.')[-1].lower()
            if file_type not in filetype_load_map:
                return {
                    "response": "Unsupport file type",
                    "status": 500,
                    "time": Time
                }
            else:
                if args['name'] == None or args['name'] == '':
                    return {
                        "response": "projectName is empty",
                        "status": 500,
                        "time": Time
                    }
                projectName = str(args['name']).strip()
                loadDocsFiles(file, size, projectName, metadatas, raw_texts, file_type)

        return {
            "params": params,
            "response": {
                "raw_texts": raw_texts,
                "metadatas": metadatas
            },
            "status": 200,
            "time": Time
        }
        ...

def loadDocsFiles(file, size, projectName, metadatas, raw_texts, file_type = None):
    global embeddings
    if file_type == None:
        file_type = file.filename.split('.')[-1].lower()
    fileName = ''
    if isinstance(file, str):
        fileName = file
    else:
        fileBytes = file.file.read()
        fileName = config.cache_path + '/' + str(uuid.uuid4()) + '.' + file_type
        fout = open(fileName, 'wb')
        fout.write(fileBytes)
        fout.close()
    documents = {}
    loader = filetype_load_map[file_type](fileName)
    text_splitter = CharacterTextSplitter(chunk_size=size, chunk_overlap=0, add_start_index=True)
    documents = loader.load_and_split(text_splitter)
    # texts = text_splitter.split_documents(documents)
    loadEmbeddings()
    path = config.wiki_path + '/' + projectName + '/' + str(size)

    for t in documents:
        start_index = t.metadata['start_index']
        page_num = t.metadata['page'] if 'page' in t.metadata else 0
        metadatas.append({'source': f'{fileName}:P{page_num}_O{start_index}'})
        raw_texts.append(t.page_content)
        # 检测知识库是否已经有相关信息
        isNew = True
        db = None
        if os.path.exists(path):
            db = FAISS.load_local(path, embeddings)
            docs = db.similarity_search(t.page_content, k=1)
            if len(docs) > 0:
                for tt in docs:
                    if tt.page_content == t.page_content:
                        isNew = False
        if isNew == True:
            docs = [t]
            if db != None: del db
            db = None
            if os.path.exists(path) == False:
                db = FAISS.from_documents(docs, embeddings)
            else:
                db = FAISS.load_local(path, embeddings)
                db.add_documents(docs)
            db.save_local(path)
    return {
        'metadatas': metadatas,
        'raw_texts': raw_texts
    }

def clear():
    os_name = platform.system()
    clear_command = 'cls' if os_name == 'Windows' else 'clear'
    os.system(clear_command)

if __name__ == '__main__':
    clear()
    uvicorn.run(app='main:app', host='0.0.0.0', port=config.port, workers=2)